#include	"../include/Socket.h"
#include	"../include/Logger.h"

#if defined(_WIN32)
#	include		<WinSock2.h>
#	include		<WS2tcpip.h>
#	pragma		comment(lib, "ws2_32.lib")
#else
#	include		<arpa/inet.h>
#	include		<fcntl.h>
#	include		<netinet/in.h>
#	include		<sys/epoll.h>
#	include		<sys/socket.h>
#	include		<sys/types.h>
#	include		<unistd.h>
#	include		<cstring>
#endif

Socket::Socket() {
#if defined(_WIN32)
	WSADATA data;
	if (WSAStartup(MAKEWORD(2, 2), &data)) throw std::runtime_error("WSAStartup failed!!!");
	ctx = (int)INVALID_SOCKET;
#else
	ctx = -1;
#endif
}

Socket::~Socket() {
	Close();
#if defined(_WIN32)
	WSACleanup();
#endif
}

bool Socket::Connect(const std::string & ip, int port) {
#if defined(_WIN32)
	LOG_ASSERTR(ctx == (int)INVALID_SOCKET, false, "Socket::Connect failed. Already connected!");
#else
	LOG_ASSERTR(ctx < 0, false, "Socket::Connect failed. Already connected!");
#endif

	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(addr));

	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);

	LOG_ASSERTR(inet_pton(AF_INET, ip.c_str(), &addr.sin_addr) == 1, false, "Socket::Connect failed. Bad remote address : %s!", ip.c_str());

#if defined(_WIN32)
	SOCKET fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	LOG_ASSERTR(fd != INVALID_SOCKET, false, "Socket::Connect failed due to create socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)");

	u_long flags = 1;
	if (ioctlsocket(fd, FIONBIO, &flags) != 0) {
		closesocket(fd);
		LOG_ERR("Socket::Connect failed due to ioctlsocket(fd, FIONBIO, 1)");
		return false;
	}
#else
	int fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	LOG_ASSERTR(fd >= 0, false, "Socket::Connect failed due to create socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)");

	int flags = fcntl(fd, F_GETFL, 0);
	if (fcntl(fd, F_SETFL, O_NONBLOCK | flags) < 0) {
		close(fd);
		LOG_ERR("Socket::Connect failed due to fcntl(fd, F_SETFL, O_NONBLOCK | flags)");
		return false;
	}
#endif

	if (connect(fd, (sockaddr *)&addr, sizeof(sockaddr)) < 0) {
#if defined(_WIN32)
		int err = WSAGetLastError();
		if (err != WSAEWOULDBLOCK) {
			closesocket(fd);
			LOG_ERR("Socket::Connect failed with errcode : %d", err);
			return false;
		}

		fd_set wait_set;
		struct timeval wait_time;

		wait_time.tv_sec = 1;
		wait_time.tv_usec = 0;

		FD_ZERO(&wait_set);
		FD_SET(fd, &wait_set);

		if (select(0, 0, &wait_set, 0, &wait_time) <= 0) {
			closesocket(fd);
			LOG_ERR("Socket::Connect failed due to timeout");
			return false;
		}
#else
		if (errno != EINPROGRESS) {
			close(fd);
			LOG_ERR("Socket::Connect failed with errcode : %d", errno);
			return false;
		}

		int wait_fd = epoll_create(1);
		LOG_ASSERTR(wait_fd >= 0, false, "Socket::Connect failed due to epoll_create(1) for async waiting");

		epoll_event ev;
		ev.events = EPOLLET | EPOLLOUT | EPOLLRDHUP;
		epoll_ctl(wait_fd, EPOLL_CTL_ADD, fd, &ev);

		int count = epoll_wait(wait_fd, &ev, 1, 1000);
		close(wait_fd);
		if (count < 0 || ev.events & EPOLLRDHUP) {
			close(fd);
			LOG_ERR("Socket::Connect failed due to timeout");
			return false;
		}
#endif
	}

	ctx = (int)fd;
	return true;
}

bool Socket::Listen(const std::string & ip, int port) {
#if defined(_WIN32)
	LOG_ASSERTR(ctx == (int)INVALID_SOCKET, false, "Socket::Listen failed. Already connected!");
#else
	LOG_ASSERTR(ctx < 0, false, "Socket::Listen failed. Already connected!");
#endif

	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(addr));

	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);

	LOG_ASSERTR(inet_pton(AF_INET, ip.c_str(), &addr.sin_addr) == 1, false, "Socket::Listen failed. Bad listen address : %s!", ip.c_str());

#if defined(_WIN32)
	SOCKET fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	LOG_ASSERTR(fd != INVALID_SOCKET, false, "Socket::Listen failed due to create socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)");

	u_long flags = 1;
	if (ioctlsocket(fd, FIONBIO, &flags) != 0) {
		closesocket(fd);
		LOG_ERR("Socket::Listen failed due to ioctlsocket(fd, FIONBIO, 1)");
		return false;
	}

	if (::bind(fd, (sockaddr *)&addr, sizeof(sockaddr)) < 0 || ::listen(fd, FD_SETSIZE) < 0) {
		closesocket(fd);
		LOG_ERR("Socket::Listen failed due to bind() or listen()");
		return false;
	}
#else
	int fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	LOG_ASSERTR(fd >= 0, false, "Socket::Listen failed due to create socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)");

	int flags = fcntl(fd, F_GETFL, 0);
	if (fcntl(fd, F_SETFL, O_NONBLOCK | flags) < 0) {
		close(fd);
		LOG_ERR("Socket::Listen failed due to fcntl(fd, F_SETFL, O_NONBLOCK | flags)");
		return false;
	}

	int reuse = 1;
	setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse));

	if (::bind(fd, (sockaddr *)&addr, sizeof(sockaddr)) < 0 || ::listen(fd, 1024) < 0) {
		close(fd);
		LOG_ERR("Socket::Listen failed due to bind() or listen()");
		return false;
	}
#endif	

	ctx = (int)fd;
	return true;
}

Socket * Socket::Accept(uint32_t & ip) {
#if defined(_WIN32)
	if (ctx == (int)INVALID_SOCKET) return nullptr;

	sockaddr_in addr = { 0 };
	int addr_size = sizeof(addr);
	SOCKET ctx_accept = accept((SOCKET)ctx, (sockaddr *)&addr, &addr_size);
	if (ctx_accept == INVALID_SOCKET) return nullptr;

	u_long flags = 1;
	if (ioctlsocket(ctx_accept, FIONBIO, &flags) != 0) {
		closesocket(ctx_accept);
		return nullptr;
	}

	Socket * socket = new Socket;
	socket->ctx = (int)ctx_accept;
	ip = addr.sin_addr.s_addr;
	return socket;
#else
	if (ctx < 0) return nullptr;

	sockaddr_in addr = { 0 };
	socklen_t addr_size = sizeof(addr);
	int fd = accept(ctx, (sockaddr *)&addr, &addr_size);
	if (fd < 0) return nullptr;

	int flags = fcntl(fd, F_GETFL, 0);
	if (fcntl(fd, F_SETFL, O_NONBLOCK | flags) < 0) {
		close(fd);
		return nullptr;
	}

	Socket * socket = new Socket;
	socket->ctx = fd;
	ip = addr.sin_addr.s_addr;
	return socket;
#endif	
}

int Socket::Recv(char * buf, size_t size) {
#if defined(_WIN32)
	if (ctx == (int)INVALID_SOCKET) return -1;
	int recved = recv((SOCKET)ctx, buf, (int)size, 0);
	if (recved > 0) return recved;
	if (recved < 0 && WSAGetLastError() == WSAEWOULDBLOCK) return 0;
	return recved == 0 ? -1 : -2;
#else
	if (ctx < 0) return -1;
	int recved = (int)recv(ctx, buf, (int)size, MSG_DONTWAIT);
	if (recved > 0) return recved;
	if (recved < 0 && errno == EAGAIN) return 0;
	return recved == 0 ? -1 : -2;
#endif	
}

bool Socket::Send(const char * data, size_t size) {
#if defined(_WIN32)
	if (ctx == (int)INVALID_SOCKET) return false;

	char * ptr = (char *)data;
	int sent = 0;
	int left = (int)size;

	while (true) {
		sent = send((SOCKET)ctx, ptr, left, 0);
		if (sent < 0) {
			if (WSAGetLastError() == WSAEWOULDBLOCK) {
				Sleep(1);
			} else {
				return false;
			}
		} else if (sent < left) {
			left -= sent;
			ptr += sent;
		} else if (sent == left) {
			return true;
		} else {
			return left == 0;
		}
	}
#else
	if (ctx < 0) return false;

	char * ptr = (char *)data;
	int sent = 0;
	int left = (int)size;

	while (true) {
		sent = (int)send(ctx, ptr, left, MSG_DONTWAIT);
		if (sent < 0) {
			if (errno == EAGAIN) {
				usleep(1000);
			} else {
				return false;
			}
		} else if (sent < left) {
			left -= sent;
			ptr += sent;
		} else if (sent == left) {
			return true;
		} else {
			return left == 0;
		}
	}
#endif
}

void Socket::Close() {
#if defined(_WIN32)
	if (ctx == (int)INVALID_SOCKET) return;
	closesocket((SOCKET)ctx);
	ctx = (int)INVALID_SOCKET;
#else
	if (ctx < 0) return;
	close(ctx);
	ctx = -1;
#endif
}